import torch


def rand_bbox(batch_size, stride, bbox_classes_num):
    bbox_tmp = torch.rand(batch_size, stride, stride, bbox_classes_num)
    bbox_tmp[:, :, :, 2:5] = 0
    bbox_tmp[:, :, :, 7:] = 0
    bbox_tmp[:, :, :, [2, 3, 7, 8]] += torch.tensor([0.2, 0.15, 0.15, 0.2])
    return bbox_tmp


def bbox_iou(box1, box2):
    box1 = torch.clamp(box1, min=0, max=1)
    box2 = torch.clamp(box2, min=0, max=1)
    box2 = torch.repeat_interleave(box2.unsqueeze(0), box1.shape[0], dim=0)
    box_area = lambda box: box[:, 2] * box[:, 3]
    box1_area = box_area(box1)
    box2_area = box_area(box2)
    x1 = torch.max(box1[:, 0] - box1[:, 2] / 2, box2[:, 0] - box2[:, 2] / 2)
    y1 = torch.max(box1[:, 1] - box1[:, 3] / 2, box2[:, 1] - box2[:, 3] / 2)
    x2 = torch.min(box1[:, 0] + box1[:, 2] / 2, box2[:, 0] + box2[:, 2] / 2)
    y2 = torch.min(box1[:, 1] + box1[:, 3] / 2, box2[:, 1] + box2[:, 3] / 2)
    inner = (x1 - x2) * (y1 - y2)
    inner = torch.clamp(inner, min=0)
    union = box1_area + box2_area - inner
    return inner / union


def dst_bbox(gt, cls, num_bbox, num_classes=2):
    stride = 7
    stride_width = 1 / stride
    stride_height = 1 / stride
    bbox_classes_num = num_bbox * 5 + num_classes + 1
    grid_cells = rand_bbox(cls.shape[0], stride, bbox_classes_num)
    x, y, w, h = gt[:, 0], gt[:, 1], gt[:, 2], gt[:, 3]
    row = (x // stride_width).int()
    col = (y // stride_height).int()
    for i in range(gt.shape[0]):
        grid_cells[i, row, col, num_bbox * 5:] = cls
    grid_cells = grid_cells.view(cls.shape[0], -1, bbox_classes_num)
    p_c = torch.max(grid_cells[:,:,num_bbox * 5:],dim=-1).values
    for i, grid_cell in enumerate(grid_cells):
        grid_cell[:, 4] = bbox_iou(grid_cell[:, 0:4], gt[i]) * p_c[i]
        grid_cell[:, 9] = bbox_iou(grid_cell[:, 5:9], gt[i]) * p_c[i]

    grid_cells.contiguous()
    return grid_cells.view(-1, stride, stride, bbox_classes_num)


# if __name__ == "__main__":
#     ground_truths = torch.rand(2, 4)
#     result = dst_bbox(ground_truths, 2, 2)
#     print(result)
